import {
    WebSocketGateway,
    WebSocketServer,
    SubscribeMessage,
    OnGatewayConnection,
    OnGatewayDisconnect,
    ConnectedSocket,
    MessageBody
} from '@nestjs/websockets'
import { Server, Socket } from 'socket.io'
import { ComfyService } from './comfy.service'
import { Logger } from '@nestjs/common'

@WebSocketGateway({
    cors: {
        origin: '*'
    }
})
export class ComfyGateway implements OnGatewayConnection, OnGatewayDisconnect {
    private readonly logger = new Logger(ComfyGateway.name)
    private readonly clients = new Map<string, Socket>()

    @WebSocketServer()
    server: Server

    constructor(private readonly comfyService: ComfyService) { }

    handleConnection(client: Socket) {
        this.logger.log(`Client connected: ${client.id}`)
        this.clients.set(client.id, client)
    }

    handleDisconnect(client: Socket) {
        this.logger.log(`Client disconnected: ${client.id}`)
        this.clients.delete(client.id)
    }

    @SubscribeMessage('generate')
    async handleGenerate(
        @ConnectedSocket() client: Socket,
        @MessageBody() data: {
            type: string
            parameters: {
                prompt: string
                negativePrompt: string
                width: number
                height: number
                seed: number
            }
        }
    ) {
        try {
            // 开始生成任务
            const taskId = await this.comfyService.startGeneration(data.parameters)

            // 监听 ComfyUI 的进度
            this.comfyService.onProgress(taskId, (progress) => {
                client.emit('message', {
                    type: 'progress',
                    progress: progress.value,
                    message: progress.message
                })
            })

            // 监听生成完成
            this.comfyService.onComplete(taskId, (result) => {
                client.emit('message', {
                    type: 'complete',
                    imageUrl: result.imageUrl
                })
            })

            // 监听错误
            this.comfyService.onError(taskId, (error) => {
                client.emit('message', {
                    type: 'error',
                    message: error.message
                })
            })

        } catch (error) {
            this.logger.error('Generation error:', error)
            client.emit('message', {
                type: 'error',
                message: '生成失败'
            })
        }
    }

    // 广播消息给所有客户端
    broadcastMessage(message: any) {
        this.server.emit('message', message)
    }

    // 发送消息给特定客户端
    sendToClient(clientId: string, message: any) {
        const client = this.clients.get(clientId)
        if (client) {
            client.emit('message', message)
        }
    }
} 